'''
历史消息汇总与对话记忆管理示例

在长对话场景中，大模型的上下文窗口有限（如4K、8K、16K tokens），
当对话历史超过限制时，我们需要对历史消息进行智能汇总，保留关键信息。

核心概念：
1. 对话历史管理：存储和检索用户与AI的交互记录
2. 消息汇总：将冗长的对话历史压缩成简洁的摘要
3. 记忆管理：平衡细节保留和上下文长度限制
'''
import os

from dotenv import load_dotenv
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnablePassthrough, RunnableWithMessageHistory
from langchain_openai import ChatOpenAI

# 自动加载 .env 文件中的环境变量
load_dotenv()

# 初始化大语言模型
llm = ChatOpenAI(api_key=os.getenv("MODELSCOPE_API_KEY"), openai_api_base=os.getenv("MODELSCOPE_API_BASE"), model="deepseek-ai/DeepSeek-R1-0528", streaming=True, temperature=0.7)

# ========== 模拟构建历史消息记录 ==========
'''
ChatMessageHistory 用于存储对话历史，包含：
- user_message: 用户发送的消息
- ai_message: AI回复的消息

在实际应用中，这些历史消息通常会存储在数据库或缓存中
'''
temp_chat_history = ChatMessageHistory()

# 添加模拟对话历史，模拟一个完整的对话流程
temp_chat_history.add_user_message("hello,我叫未来")
temp_chat_history.add_ai_message("你好")
temp_chat_history.add_user_message("今天天气怎么样")
temp_chat_history.add_ai_message("今天天气晴朗，适合出去玩")
temp_chat_history.add_user_message("那我准备出门")
temp_chat_history.add_ai_message("好的，祝你玩的愉快")

print("=== 初始对话历史 ===")
for i, message in enumerate(temp_chat_history.messages):
    print(f"{i + 1}. {message.type}: {message.content}")

# ========== 构建基础对话链 ==========
'''
ChatPromptTemplate 定义对话的模板结构：
- system: 系统角色设定
- MessagesPlaceholder: 历史消息的占位符，运行时会被实际历史替换
- human: 当前用户输入
'''
prompt = ChatPromptTemplate.from_messages([
    ("system", "你是一个乐于助人的助手。尽力回答所有问题，提供的聊天历史包括与您交谈的用户的记录。"),
    MessagesPlaceholder(variable_name="chat_history"),  # 历史消息注入点
    ("human", "{input}")  # 当前用户输入
])

# 构建处理链：提示模板 -> 大模型
chain = prompt | llm

# ========== 添加历史消息管理 ==========
'''
RunnableWithMessageHistory 的核心作用：
1. 自动管理对话历史
2. 根据session_id隔离不同对话
3. 自动将历史消息注入到提示词中
'''
chain_with_message_history = RunnableWithMessageHistory(
    chain,  # 基础处理链
    lambda session_id: temp_chat_history,  # 历史消息提供函数（根据session_id返回对应的历史）
    input_messages_key="input",  # 用户输入在输入字典中的key
    history_messages_key="chat_history",  # 历史消息在提示词模板中的key
)


# ========== 消息汇总功能 ==========
def summarize_messages(chain_input):
    """
    智能汇总历史消息，将长对话压缩成简洁摘要

    应用场景：
    - 对话历史超过模型上下文限制时
    - 需要长期记忆但不想存储所有细节时
    - 跨会话保持关键信息时

    工作流程：
    1. 检查是否有历史消息需要汇总
    2. 使用大模型生成摘要
    3. 清空原始历史，只保留摘要
    4. 返回处理状态
    """
    stored_messages = temp_chat_history.messages

    # 检查是否有历史消息需要处理
    if len(stored_messages) == 0:
        print("📭 没有历史消息需要汇总")
        return False

    print(f"📚 开始汇总 {len(stored_messages)} 条历史消息...")

    # 构建摘要生成提示词
    summarization_prompt = ChatPromptTemplate.from_messages([
        MessagesPlaceholder(variable_name="chat_history"),  # 传入完整历史
        ("user", "将上述聊天消息浓缩成一条摘要消息，尽可能包含多个具体细节。包括用户的重要信息、讨论的主题和关键决定。")
    ])

    # 创建摘要生成链
    summarization_chain = summarization_prompt | llm

    # 生成摘要
    summary_message = summarization_chain.invoke({"chat_history": stored_messages})

    print("✅ 摘要生成完成:")
    print(f"   原始消息数: {len(stored_messages)}")
    print(f"   摘要内容: {summary_message.content}")

    # 清空原始历史，用摘要替换
    temp_chat_history.clear()
    temp_chat_history.add_message(summary_message)

    print("🔄 历史记录已更新：原始对话 → 智能摘要")
    return True


# ========== 构建带汇总功能的完整链 ==========
'''
chain_with_summarization 的工作流程：
1. RunnablePassthrough.assign(messages_summarized=summarize_messages)
   - 先执行 summarize_messages 函数
   - 将返回值赋给 messages_summarized 变量
   - 保持原始输入传递

2. | chain_with_message_history
   - 将处理后的输入传递给带历史管理的对话链
   - 生成最终回复

这种设计确保在每次对话前都会检查是否需要汇总历史
'''
chain_with_summarization = (RunnablePassthrough.assign(messages_summarized=summarize_messages) | chain_with_message_history)

# ========== 测试演示 ==========
print("\n" + "=" * 60)
print("🚀 开始测试历史消息汇总功能")
print("=" * 60)

# 模拟用户询问需要历史记忆的问题
test_questions = [
    "我是谁？下午在干嘛？准备出门吗？",
    # 可以继续测试更多问题来观察历史管理效果
]

for i, question in enumerate(test_questions, 1):
    print(f"\n📝 测试问题 {i}: {question}")

    # 调用带汇总功能的对话链
    response = chain_with_summarization.invoke(
        {"input": question},
        {"configurable": {"session_id": "lqq"}}  # session_id用于区分不同对话会话
    )

    print(f"🤖 AI回复: {response.content}")

    # 显示当前的历史状态
    print("\n📋 当前历史记录状态:")
    for j, message in enumerate(temp_chat_history.messages):
        print(f"   {j + 1}. {message.type}: {message.content}")


# ========== 进阶演示：模拟长对话场景 ==========
def simulate_long_conversation():
    """
    模拟长对话场景，展示历史汇总的实际价值
    """
    print("\n" + "=" * 60)
    print("🔬 进阶演示：模拟长对话场景")
    print("=" * 60)

    # 清空历史，重新开始
    temp_chat_history.clear()

    # 模拟多轮对话
    conversations = [
        ("你好，我叫张三，是一名软件工程师", "助手回复..."),
        ("我最近在学习机器学习，特别是深度学习", "助手回复..."),
        ("我打算用Python和TensorFlow来构建一个图像识别项目", "助手回复..."),
        ("项目需要识别猫和狗的照片，准确率要达到95%以上", "助手回复..."),
        ("我计划使用CNN卷积神经网络架构", "助手回复..."),
        ("数据集准备使用Kaggle上的公开数据集", "助手回复..."),
    ]

    # 添加多轮对话到历史
    for user_msg, ai_msg in conversations:
        temp_chat_history.add_user_message(user_msg)
        temp_chat_history.add_ai_message(ai_msg)

    print(f"📖 创建了 {len(temp_chat_history.messages)} 条对话历史")

    # 触发历史汇总
    summarize_messages({})

    # 测试汇总后的记忆保持
    test_question = "我之前说的项目是什么？要用什么技术？目标是什么？"
    print(f"\n❓ 测试问题: {test_question}")

    response = chain_with_summarization.invoke(
        {"input": test_question},
        {"configurable": {"session_id": "test_session"}}
    )

    print(f"🤖 AI回复: {response.content}")

# 取消注释运行进阶演示
# simulate_long_conversation()
